# 导入库
import numpy as np
from joblib import load

class Classification(object):
    def __init__(self):
        self.knn = load('knn_model.joblib')
        self.scaler = load('scaler.joblib')
        self.le = load('label_encoder.joblib')
    
    def get_predictor_category(self, X):
        """
        获取预测的分类
        """
        # 转化为数组并保持特征顺序
        features_order = ['eps', 'total_revenue_ps', 'undist_profit_ps', 'gross_margin', 'fcff', 'fcfe', 'tangible_asset', 'bps', 
                         'grossprofit_margin', 'npta']
        new_value = np.array([X[feature] for feature in features_order])
        
        # 标准化新数据
        new_scaled = self.scaler.transform([new_value])
        
        # 预测分类
        predicted_label = self.knn.predict(new_scaled)
        
        # 解码分类标签
        predicted_category = self.le.inverse_transform(predicted_label)
        return predicted_category[0]

if __name__ == '__main__':
    ci = Classification()
    new_data = {
        'eps': 1.59,
        'total_revenue_ps': 19.4516,
        'undist_profit_ps': 7.0876,
        'gross_margin': 4063940000,
        'fcff': 27075800,
        'fcfe': 44465800,
        'tangible_asset': 9468120000,
        'bps': 11.5354,
        'grossprofit_margin': 23.4565,
        'npta': 8.3336
    }
    
    # 预测分类
    predicted_category = ci.get_predictor_category(new_data)
    print(f"预测的分类: {predicted_category}")